import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def plot_rigid_transformations(matrix1_flat, matrix2_flat, axis_length=1.0):
    """ 可视化三个刚体变换矩阵（包含RGB单位矩阵）的三维坐标系 """
    def to_homogeneous(flat_array):
        arr = np.array(flat_array).reshape(3, 4)
        return np.vstack([arr, [0, 0, 0, 1]])

    # 转换为齐次矩阵
    T0 = np.eye(4)  # 单位矩阵
    T1 = to_homogeneous(matrix1_flat)
    T2 = to_homogeneous(matrix2_flat)

    # 提取旋转矩阵和平移向量
    R0, t0 = T0[:3, :3], T0[:3, 3]
    R1, t1 = T1[:3, :3], T1[:3, 3]
    R2, t2 = T2[:3, :3], T2[:3, 3]

    fig = plt.figure(figsize=(14, 10))
    ax = fig.add_subplot(111, projection='3d')
    ax.dist = 12  # 初始观察距离

    # 滚轮缩放功能
    def on_scroll(event):
        if event.inaxes != ax: return
        ax.dist *= 0.9 if event.step > 0 else 1.1
        ax.dist = np.clip(ax.dist, 5, 35)
        fig.canvas.draw_idle()
    fig.canvas.mpl_connect('scroll_event', on_scroll)

    def create_segment_with_marker(start, direction, total_length, color, marker_type):
        """ 创建带末端标记的彩色线段 """
        # 公共参数
        arrow_ratio = 0.2
        line_length = total_length * (1 - arrow_ratio)
        arrow_length = total_length * arrow_ratio
        dir_normalized = direction / np.linalg.norm(direction)
        
        # 绘制线段
        line_end = start + dir_normalized * line_length
        ax.plot([start[0], line_end[0]], 
                [start[1], line_end[1]],
                [start[2], line_end[2]],
                color=color, lw=2.5, alpha=0.9 if marker_type!='unit' else 0.8)

        # 生成标记几何体
        tip = line_end + dir_normalized * arrow_length
        if marker_type == 'cone':
            return create_cone_marker(line_end, dir_normalized, arrow_length, color)
        elif marker_type == 'pyramid':
            return create_pyramid_marker(line_end, dir_normalized, arrow_length, color)
        elif marker_type == 'unit':
            return create_cylinder_marker(line_end, dir_normalized, arrow_length, color)

    def create_cone_marker(base, direction, length, color):
        """ 生成圆锥标记 """
        radius = length * 0.3
        if np.abs(direction[0]) < 0.1 and np.abs(direction[1]) < 0.1:
            u = np.array([1.0, 0.0, 0.0])
            v = np.array([0.0, 1.0, 0.0])
        else:
            u = np.array([direction[1], -direction[0], 0.0])
            u /= np.linalg.norm(u)
            v = np.cross(direction, u)
        theta = np.linspace(0, 2*np.pi, 20)
        points = [base + radius*(u*np.cos(t) + v*np.sin(t)) for t in theta]
        tip = base + direction*length
        return Poly3DCollection(
            [[points[i], points[(i+1)%20], tip] for i in range(20)],
            facecolor=color, edgecolor='none', alpha=0.9
        )

    def create_pyramid_marker(base, direction, length, color):
        """ 生成方锥标记 """
        width = length * 0.4
        if np.abs(direction[0]) < 0.1 and np.abs(direction[1]) < 0.1:
            u = np.array([width, 0, 0])
            v = np.array([0, width, 0])
        else:
            u = np.array([direction[1], -direction[0], 0])
            u = width * u / np.linalg.norm(u)
            v = np.cross(direction, u)
            v = width * v / np.linalg.norm(v)
        points = [
            base - u - v,
            base + u - v,
            base + u + v,
            base - u + v
        ]
        tip = base + direction*length
        return Poly3DCollection(
            [[points[0], points[1], tip],
             [points[1], points[2], tip],
             [points[2], points[3], tip],
             [points[3], points[0], tip],
             points],
            facecolor=color, edgecolor='none', alpha=0.7
        )

    def create_cylinder_marker(base, direction, length, color):
        """ 生成圆柱标记（用于单位坐标系） """
        radius = length * 0.2
        height = length * 1.2
        if np.abs(direction[0]) < 0.1 and np.abs(direction[1]) < 0.1:
            u = np.array([1.0, 0.0, 0.0])
            v = np.array([0.0, 1.0, 0.0])
        else:
            u = np.array([direction[1], -direction[0], 0.0])
            u /= np.linalg.norm(u)
            v = np.cross(direction, u)
        theta = np.linspace(0, 2*np.pi, 12)
        triangles = []
        for i in range(len(theta)-1):
            p1 = base + radius*(u*np.cos(theta[i]) + v*np.sin(theta[i]))
            p2 = base + radius*(u*np.cos(theta[i+1]) + v*np.sin(theta[i+1]))
            p3 = p2 + direction*height
            p4 = p1 + direction*height
            triangles.extend([[p1, p2, p3, p4], [p1, p4, p1+height*direction]])
        return Poly3DCollection(
            triangles,
            facecolor=color, edgecolor='none', alpha=0.8
        )

    # 颜色定义 (单位坐标系使用标准RGB)
    colors = {
        'unit':   [(1.0, 0.0, 0.0), (0.0, 0.8, 0.0), (0.0, 0.0, 1.0)],  # 纯红、绿、蓝
        'cone':   [(1.0, 0.2, 0.2), (0.2, 0.8, 0.2), (0.2, 0.4, 1.0)],  # 变换1
        'pyramid':[(1.0, 0.6, 0.6), (0.6, 1.0, 0.6), (0.6, 0.8, 1.0)]   # 变换2
    }

    # 绘制三个坐标系
    for i, (R, t, style) in enumerate([
        (R0, t0, ('unit', 0)),  # 单位矩阵
        (R1, t1, ('cone', 0)),
        (R2, t2, ('pyramid', 0))
    ]):
        marker_type = style[0]
        for axis in range(3):
            direction = R[:, axis] * axis_length
            color = colors[marker_type][axis]
            coll = create_segment_with_marker(
                t, direction, axis_length, 
                color=color,
                marker_type=marker_type if marker_type=='unit' else style[0]
            )
            ax.add_collection3d(coll)

    # 设置坐标范围
    all_origins = np.vstack([t0, t1, t2])
    max_range = np.ptp(all_origins, axis=0).max() * 3
    mid_point = (all_origins.max(axis=0) + all_origins.min(axis=0)) * 0.5
    ax.set_xlim(mid_point[0]-max_range, mid_point[0]+max_range)
    ax.set_ylim(mid_point[1]-max_range, mid_point[1]+max_range)
    ax.set_zlim(mid_point[2]-max_range, mid_point[2]+max_range)

    # 添加标签和图例
    ax.set_xlabel('X', color=colors['unit'][0], fontsize=12, fontweight='bold')
    ax.set_ylabel('Y', color=colors['unit'][1], fontsize=12, fontweight='bold')
    ax.set_zlabel('Z', color=colors['unit'][2], fontsize=12, fontweight='bold')
    ax.set_title('Rigid Transformations with RGB Reference Frame', pad=25)

    # 创建图例
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor=colors['unit'][0], label='Reference X'),
        Patch(facecolor=colors['unit'][1], label='Reference Y'),
        Patch(facecolor=colors['unit'][2], label='Reference Z'),
        Patch(facecolor=colors['cone'][0], label='Transform 1 (Cone)'),
        Patch(facecolor=colors['pyramid'][0], label='Transform 2 (Pyramid)')
    ]
    ax.legend(handles=legend_elements, loc='upper right', fontsize=9)

    plt.tight_layout()
    plt.show()

# 测试样例
if __name__ == "__main__":
    # 变换1：绕Z轴旋转30度
    theta = np.pi/6
    transform1 = np.array([
        [np.cos(theta), -np.sin(theta), 0, 1.5],
        [np.sin(theta), np.cos(theta), 0, -0.5],
        [0, 0, 1, 0.5]
    ]).flatten()

    # 变换2：复杂变换
    phi = np.pi/4
    transform2 = np.array([
        [np.cos(phi), 0, np.sin(phi), 2.0],
        [0, 1, 0, 1.0],
        [-np.sin(phi), 0, np.cos(phi), 0.5],
    ]).flatten()

    plot_rigid_transformations(transform1, transform2, axis_length=1.5)